import os
from convert_model import convert_model

def convert_all_models(model_dir):
    """转换目录中的所有.th模型"""
    for filename in os.listdir(model_dir):
        if filename.endswith('.th'):
            input_path = os.path.join(model_dir, filename)
            try:
                print(f"\nConverting {filename}...")
                convert_model(input_path)
            except Exception as e:
                print(f"Failed to convert {filename}: {str(e)}")

if __name__ == "__main__":
    model_dir = "model_files"  # 模型文件夹路径
    convert_all_models(model_dir) 